import math
import torch as pt
import numpy as np
from vcpy.m3d import normalize

def generate_sphere_cams(ref_rot, ref_t,
  total_frame=120, spin_radius=0.1, total_spin=1):

  def dcm_to_4x4(r,t):
    camera_matrix = pt.eye(4, dtype=pt.float32)
    camera_matrix[:3,:3] = r
    if t.dim() > 1:
      camera_matrix[:3,3:4] = t
    else:
      camera_matrix[:3,3] = t
    return camera_matrix

  ref_cam = dcm_to_4x4(ref_rot, ref_t)
  spin_speed = 2*math.pi / total_frame * total_spin
  render_poses = {}
  for i in range(total_frame):
    anim_time = spin_speed * i
    x = math.cos(anim_time) * spin_radius
    y = math.sin(anim_time) * spin_radius
    euler_4x4 = dcm_to_4x4(pt.eye(3), pt.tensor([x,y,0.0]))
    output = euler_4x4 @ ref_cam
    output = output.to(pt.float32)
    render_poses[i] = {'rot': output[:3, :3], 't': output[:3, 3]}
  return render_poses

def __get_cam_frame(origin, target, up):
  F = target[:3] - origin[:3]
  f = normalize(F)
  U = normalize(up[:3])
  s = normalize(np.cross(U,f))
  u = np.cross(f,s)
  M = np.identity(4)
  M[:3,:3] = np.c_[s,u,f]
  M[:3,3] = origin[:3]
  return M

# lookat fixed point
def generate_spiral_cams(radius_xyz, lookat_pos, up, zrate=0.5, rots=2, N=120):
  render_poses = {}

  for i, theta_2 in enumerate(np.linspace(0, (2 * np.pi * rots) ** 2, N+1)[:-1]):
    theta = theta_2 ** 0.5
    cam_pos = np.array([np.cos(theta) * i / N, -np.sin(theta) * i / N, -np.sin(theta*zrate)]) * radius_xyz
    mat = pt.from_numpy(__get_cam_frame(cam_pos, lookat_pos, up)).to(pt.float)
    render_poses[i] = {'rot': mat[:3, :3], 't': mat[:3, 3]}
  return render_poses

# lookat inside-out direction
def generate_spiral_cams2(radius_xyz, inv_lookat_pos, up, zrate=0.5, rots=2, N=120):
  render_poses = {}

  for i, theta_2 in enumerate(np.linspace(0, (2 * np.pi * rots) ** 2, N+1)[:-1]):
    theta = theta_2 ** 0.5
    cam_pos = np.array([np.cos(theta) * i / N, -np.sin(theta) * i / N, -np.sin(theta*zrate)]) * radius_xyz
    lookat_pos = 2 * cam_pos - inv_lookat_pos
    mat = pt.from_numpy(__get_cam_frame(cam_pos, lookat_pos, up)).to(pt.float)
    render_poses[i] = {'rot': mat[:3, :3], 't': mat[:3, 3]}
  return render_poses
